// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use std::{
    sync::{Arc, Once},
    thread,
    time::{Duration, Instant},
};

use tracing::{debug, trace};

use super::*;

pub fn init_test_logger() {
    static INIT_LOGGER: Once = Once::new();

    INIT_LOGGER.call_once(|| {
        let _ = tracing_subscriber::fmt()
            .with_max_level(tracing::Level::DEBUG)
            .with_thread_ids(true)
            .with_thread_names(false)
            .with_file(false)
            .with_target(false)
            .with_writer(std::io::stdout)
            .try_init();
    });
}

pub fn create_connection(
    transport: Transport,
    addr: &str,
    timeout: Duration,
) -> Result<(Endpoint, Endpoint), Error> {
    let transport = Arc::new(transport);

    let server = transport.create(addr)?;
    let client = transport.connect(addr, timeout)?;

    Ok((server, client))
}

pub fn send_message(endpoint: &mut Endpoint, message: &[u8]) -> Result<(), Error> {
    let msg_len = message.len();

    trace!("Sending message '{}'", String::from_utf8_lossy(message));

    // Write message to buffer
    let mut write_buf = endpoint.write()?;
    write_buf[..message.len()].copy_from_slice(message);

    // Submit changes
    write_buf.submit(message.len())?;
    trace!("Sent {} bytes", msg_len);

    Ok(())
}

pub fn receive_message(endpoint: &mut Endpoint, expected: &[u8]) -> Result<(), Error> {
    let msg_len = expected.len();

    trace!("Receiving message...");

    // Read message from buffer
    let read_buf = endpoint.read()?;
    let message = &read_buf[..expected.len()];

    // Validate message
    trace!("Received message '{}'", String::from_utf8_lossy(message));
    assert_eq!(message, expected, "Received message mismatch");

    // Consume read buffer
    read_buf.consume(expected.len())?;
    trace!("Received {} bytes", msg_len);

    Ok(())
}

pub fn send_raw_bytes(endpoint: &mut Endpoint, message: &[u8]) -> Result<(), Error> {
    let total_len = message.len();
    let mut bytes_sent = 0;

    while bytes_sent < total_len {
        let mut write_buf = endpoint.write()?;

        // Determine the size of the next chunk to send
        let chunk_size = std::cmp::min(
            write_buf.len(),        // Buffer capacity
            total_len - bytes_sent, // Remaining bytes
        );

        let chunk = &message[bytes_sent..bytes_sent + chunk_size];
        write_buf[..chunk.len()].copy_from_slice(chunk);

        write_buf.submit(chunk.len())?;
        bytes_sent += chunk.len();

        trace!("Sent: {}/{}, chunk: {}", bytes_sent, total_len, chunk_size);
    }

    Ok(())
}

pub fn recv_raw_bytes(endpoint: &mut Endpoint, expected: &[u8]) -> Result<Vec<u8>, Error> {
    let expected_len = expected.len();
    let mut received_data = Vec::with_capacity(expected_len);

    while received_data.len() < expected_len {
        let read_buf = endpoint.read()?;

        let received_len = read_buf.len();
        if received_len == 0 {
            break;
        }
        received_data.extend_from_slice(&read_buf);

        trace!(
            "Recv: {}/{}, chunk: {}",
            received_data.len(),
            expected_len,
            received_len
        );
        read_buf.consume(received_len)?;
    }

    Ok(received_data)
}

pub fn bidirectional_communication(
    transport: Transport,
    address: &str,
    timeout: Duration,
) -> Result<(), Error> {
    const PING: &[u8] = b"Ping";
    const PONG: &[u8] = b"Pong";

    // Initialize logger
    self::init_test_logger();

    // Create server & client
    let (mut server, mut client) = self::create_connection(transport, address, timeout)?;

    for _ in 0..100 {
        // Client send PING
        self::send_message(&mut client, PING)?;

        // Server receive PING
        self::receive_message(&mut server, PING)?;

        // Server send PONG
        self::send_message(&mut server, PONG)?;

        // Client receive PING
        self::receive_message(&mut client, PONG)?;
    }

    Ok(())
}

pub fn transfer_raw_bytes(
    transport: Transport,
    address: &str,
    data_size: usize,
    timeout: Duration,
) -> Result<(), Error> {
    // Initialize logger
    self::init_test_logger();

    // Generate test data
    debug!("Generating test data...");
    let test_data = Arc::new((0..data_size).map(|i| (i % 256) as u8).collect::<Vec<_>>());

    // Create server & client
    let (mut server, mut client) = self::create_connection(transport, address, timeout)?;

    let expected_data = test_data.clone();

    // Start sender thread
    let client_thread = {
        let test_data = test_data.clone();
        thread::Builder::new()
            .name("sender".to_string())
            .spawn(move || {
                let total_bytes = test_data.len();
                let start_time = Instant::now();

                debug!("Sending {} bytes...", total_bytes);
                self::send_raw_bytes(&mut client, &test_data).expect("Send raw bytes failed");

                let duration = start_time.elapsed();
                (total_bytes, duration)
            })
            .expect("failed to start sender thread")
    };

    // Start receiver thread
    let server_thread = {
        thread::Builder::new()
            .name("receiver".to_string())
            .spawn(move || {
                let total_bytes = expected_data.len();
                let start_time = Instant::now();

                debug!("Receiving {} bytes...", total_bytes);
                let received_data = self::recv_raw_bytes(&mut server, &expected_data)
                    .expect("Receive raw bytes failed");

                let duration = start_time.elapsed();
                assert_eq!(received_data, expected_data.as_slice(), "Data mismatch");
                (total_bytes, duration)
            })
            .expect("failed to start receiver thread")
    };

    // Get results and print stats
    let (send_bytes, send_duration) = client_thread.join().expect("sender thread panicked");
    let (recv_bytes, recv_duration) = server_thread.join().expect("receiver thread panicked");

    // Print sender stats
    let send_speed = (send_bytes as f64 / 1024.0 / 1024.0) / send_duration.as_secs_f64();
    debug!(
        "Send takes {:.3}s, speed: {:.2}MB/s",
        send_duration.as_secs_f64(),
        send_speed
    );

    // Print receiver stats
    let recv_speed = (recv_bytes as f64 / 1024.0 / 1024.0) / recv_duration.as_secs_f64();
    debug!(
        "Recv takes {:.3}s, speed: {:.2}MB/s",
        recv_duration.as_secs_f64(),
        recv_speed
    );

    Ok(())
}

use crate::sys::shmem;

#[test]
fn test_bidirectional_communication() -> Result<(), Error> {
    const CONN_TIMEOUT: Duration = Duration::from_millis(200);

    let transport = Transport::default();
    let address = shmem::unique_name();

    bidirectional_communication(transport, &address, CONN_TIMEOUT)
}

#[test]
fn test_raw_bytes_transfer() -> Result<(), Error> {
    const DATA_SIZE: usize = 64 * 1024 * 1024; // 64M
    const BUFFER_SIZE: usize = 16 * 1024; // 16K
    const CONN_TIMEOUT: Duration = Duration::from_millis(200);

    let transport = Transport::new(BUFFER_SIZE);
    let address = shmem::unique_name();

    transfer_raw_bytes(transport, &address, DATA_SIZE, CONN_TIMEOUT)
}
